import torch
import cv2
import C3D_model
import numpy as np

def center_crop(frame):
    frame=frame[8:120,30:142,:]
    return np.array(frame).astype(np.uint8)


def inference():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")#定义模型设备
    #加载数据标签
    with open("./data/labels.txt", "r") as f:
        class_names = f.readlines()
        print(class_names)
        f.close()

    #加载模型，并将模型参数加到模型中
    model = C3D_model.C3D(num_classes=101)
    checkpoint=torch.load('model_result/models/C3D_epoch-29.pth.tar')#加载模型
    model.load_state_dict(checkpoint['state_dict'])#模型参数

    #将模型放入到设备中，并设置验证模式
    model.to(device)
    model.eval()

    video="./data/video.mp4"
    cap = cv2.VideoCapture(video)
    retaining=True
    clip=[]

    while retaining:
        retaining,frame=cap.read()#视频结束retaining为假，读取视频帧
        if not retaining and frame is None:
            continue
        tmp_=center_crop(cv2.resize(frame,(171,128)))
        tmp=tmp_ -np.array([[[90.0,98.0,192.0]]])#归一化
        clip.append(tmp)#clip放所有的帧
        if len(clip)==16:
            inputs=np.array(clip).astype(np.float32)
            inputs=np.expand_dims(inputs,axis=0)#升维度(1,16,112,112,3)
            inputs=np.transpose(inputs,[0,4,1,2,3])
            inputs=torch.from_numpy(inputs)
            inputs=torch.autograd.Variable(inputs,requires_grad=False).to(device)#老的一种写法

            with torch.no_grad():
                outputs=model(inputs)

            probs=torch.nn.Softmax(dim=1)(outputs)#概率
            #torch.argmax 只返回最大值的索引。torch.max 返回一个元组，包含最大值和最大值的索引。
            label=torch.max(probs,dim=1)[1].detach().cpu().numpy()[0]

            cv2.putText(frame,class_names[label].split(' ')[-1].strip(),(20,20),
                        cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,255),1)#标签

            cv2.putText(frame, 'prob:%.4f'%probs[0][label], (20, 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1)#概率

            clip.pop(0)#pop(0) 方法移除列表的第一个元素（索引为0的元素）并返回它
            cv2.imshow('result',frame)
            cv2.waitKey(30)#停留
        cap.release()
        cv2.destroyAllWindows()



if __name__ == '__main__':
    inference()#['1 ApplyEyeMakeup\n',2 ApplyLipstick\n,...]